#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
FrankMocap PKL文件合并和轨迹数据生成的完整流水线
一键运行所有步骤
"""

import os
import sys
import subprocess
import argparse
from datetime import datetime

def run_command(cmd, description):
    """运行命令并显示进度"""
    print(f"\n{'='*60}")
    print(f"🔄 {description}")
    print(f"执行命令: {cmd}")
    print(f"{'='*60}")
    
    try:
        result = subprocess.run(cmd, shell=True, check=True, 
                              capture_output=False, text=True)
        print(f"✅ {description} - 完成")
        return True
    except subprocess.CalledProcessError as e:
        print(f"❌ {description} - 失败")
        print(f"错误: {e}")
        return False

def check_dependencies():
    """检查必要的文件是否存在"""
    required_files = [
        'analyze_pkl_structure.py',
        'merge_pkls_to_pkl.py', 
        'test_trajectory_data.py'
    ]
    
    missing_files = []
    for file in required_files:
        if not os.path.exists(file):
            missing_files.append(file)
    
    if missing_files:
        print(f"❌ 缺少必要文件: {missing_files}")
        return False
    
    print("✅ 所有必要文件都存在")
    return True

def main():
    parser = argparse.ArgumentParser(description='FrankMocap PKL合并流水线')
    parser.add_argument('--input', '-i', 
                       default='./mocap_output_full/mocap',
                       help='输入pkl文件目录')
    parser.add_argument('--output-dir', '-o',
                       default='./output_pkl',
                       help='输出目录')
    parser.add_argument('--format', '-f',
                       choices=['animation_data', 'frame_list', 'both'],
                       default='animation_data',
                       help='输出格式')
    parser.add_argument('--skip-analysis', action='store_true',
                       help='跳过pkl文件结构分析')
    parser.add_argument('--skip-visualization', action='store_true',
                       help='跳过轨迹可视化')
    
    args = parser.parse_args()
    
    # 记录开始时间
    start_time = datetime.now()
    print(f"🚀 FrankMocap PKL合并流水线启动")
    print(f"开始时间: {start_time.strftime('%Y-%m-%d %H:%M:%S')}")
    
    # 检查依赖
    if not check_dependencies():
        sys.exit(1)
    
    # 检查输入目录
    if not os.path.exists(args.input):
        print(f"❌ 输入目录不存在: {args.input}")
        sys.exit(1)
    
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    os.makedirs('./trajectory_plots', exist_ok=True)
    
    success_count = 0
    total_steps = 0
    
    # 步骤1: 分析pkl文件结构
    if not args.skip_analysis:
        total_steps += 1
        cmd = f"python analyze_pkl_structure.py -i {args.input}"
        if run_command(cmd, "步骤1: 分析PKL文件结构"):
            success_count += 1
    
    # 步骤2: 合并pkl文件
    if args.format in ['animation_data', 'both']:
        total_steps += 1
        output_file = os.path.join(args.output_dir, 'animation_with_trajectory.pkl')
        cmd = f"python merge_pkls_to_pkl.py -i {args.input} -o {output_file} -f animation_data"
        if run_command(cmd, "步骤2a: 合并PKL文件 (animation_data格式)"):
            success_count += 1
            
            # 步骤3a: 测试轨迹数据
            total_steps += 1
            cmd = f"python test_trajectory_data.py -i {output_file}"
            if run_command(cmd, "步骤3a: 分析轨迹数据"):
                success_count += 1
            
            # 步骤4a: 可视化轨迹（可选）
            if not args.skip_visualization:
                total_steps += 1
                cmd = f"python test_trajectory_data.py -i {output_file} --plot"
                if run_command(cmd, "步骤4a: 生成轨迹可视化图"):
                    success_count += 1
    
    if args.format in ['frame_list', 'both']:
        total_steps += 1
        output_file = os.path.join(args.output_dir, 'frame_list.pkl')
        cmd = f"python merge_pkls_to_pkl.py -i {args.input} -o {output_file} -f frame_list"
        if run_command(cmd, "步骤2b: 合并PKL文件 (frame_list格式)"):
            success_count += 1
            
            # 步骤3b: 测试帧列表数据
            total_steps += 1
            cmd = f"python test_trajectory_data.py -i {output_file}"
            if run_command(cmd, "步骤3b: 分析帧列表数据"):
                success_count += 1
    
    # 总结
    end_time = datetime.now()
    duration = end_time - start_time
    
    print(f"\n{'='*60}")
    print(f"🎯 流水线执行完成")
    print(f"{'='*60}")
    print(f"执行时间: {duration}")
    print(f"成功步骤: {success_count}/{total_steps}")
    
    if success_count == total_steps:
        print(f"✅ 所有步骤都成功完成!")
        print(f"\n📁 输出文件位置:")
        if args.format in ['animation_data', 'both']:
            animation_file = os.path.join(args.output_dir, 'animation_with_trajectory.pkl')
            if os.path.exists(animation_file):
                file_size = os.path.getsize(animation_file) / (1024*1024)
                print(f"  - {animation_file} ({file_size:.2f} MB)")
        if args.format in ['frame_list', 'both']:
            frame_file = os.path.join(args.output_dir, 'frame_list.pkl')
            if os.path.exists(frame_file):
                file_size = os.path.getsize(frame_file) / (1024*1024)
                print(f"  - {frame_file} ({file_size:.2f} MB)")
        
        if not args.skip_visualization and os.path.exists('./trajectory_plots'):
            print(f"  - 轨迹图: ./trajectory_plots/")
        
        print(f"\n🎉 现在你可以使用这些文件进行动画处理了!")
    else:
        print(f"⚠️  有 {total_steps - success_count} 个步骤失败")
        print(f"请检查上面的错误信息")

if __name__ == "__main__":
    main()